"""
Statistics and Plotting Utilities


- 
- 
- IoU
- ROI
- Before/After
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
from PIL import Image
from typing import List, Dict, Tuple, Optional
import seaborn as sns
import logging

logger = logging.getLogger(__name__)

# 
sns.set_style("whitegrid")
plt.rcParams['font.size'] = 10


class StatisticsPlotter:
    """"""

    @staticmethod
    def plot_class_distribution(
        detections: List[Dict],
        class_names: List[str],
        figsize: Tuple[int, int] = (10, 6)
    ) -> Image.Image:
        """
        

        Args:
            detections:  [{'class': str, 'confidence': float, ...}, ...]
            class_names: 
            figsize: 

        Returns:
            PIL.Image
        """
        # 
        class_counts = {name: 0 for name in class_names}
        for det in detections:
            cls = det['class']
            if cls in class_counts:
                class_counts[cls] += 1

        # 
        classes = list(class_counts.keys())
        counts = list(class_counts.values())

        # 
        fig, ax = plt.subplots(figsize=figsize)

        # 
        bars = ax.bar(classes, counts, color='steelblue', alpha=0.7, edgecolor='black')

        # 
        for bar in bars:
            height = bar.get_height()
            if height > 0:
                ax.text(
                    bar.get_x() + bar.get_width() / 2.,
                    height,
                    f'{int(height)}',
                    ha='center',
                    va='bottom',
                    fontsize=11,
                    fontweight='bold'
                )

        ax.set_xlabel('Class', fontsize=12, fontweight='bold')
        ax.set_ylabel('Count', fontsize=12, fontweight='bold')
        ax.set_title(f'Class Distribution (Total: {len(detections)} detections)',
                     fontsize=14, fontweight='bold')
        ax.set_ylim(0, max(counts) * 1.2 if counts else 1)

        # x
        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()

        # PIL Image
        img = StatisticsPlotter._fig_to_image(fig)
        plt.close(fig)

        return img

    @staticmethod
    def plot_confidence_distribution(
        detections: List[Dict],
        bins: int = 20,
        figsize: Tuple[int, int] = (10, 6)
    ) -> Image.Image:
        """
        

        Args:
            detections: 
            bins: bin
            figsize: 

        Returns:
            PIL.Image
        """
        # 
        confidences = [det['confidence'] for det in detections]

        if len(confidences) == 0:
            # 
            fig, ax = plt.subplots(figsize=figsize)
            ax.text(0.5, 0.5, 'No detections',
                   ha='center', va='center', fontsize=16)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            img = StatisticsPlotter._fig_to_image(fig)
            plt.close(fig)
            return img

        # 
        fig, ax = plt.subplots(figsize=figsize)

        # 
        n, bins_edges, patches = ax.hist(
            confidences,
            bins=bins,
            range=(0, 1),
            color='coral',
            alpha=0.7,
            edgecolor='black'
        )

        # 
        mean_conf = np.mean(confidences)
        median_conf = np.median(confidences)

        # 
        ax.axvline(mean_conf, color='red', linestyle='--',
                  linewidth=2, label=f'Mean: {mean_conf:.3f}')
        ax.axvline(median_conf, color='green', linestyle='--',
                  linewidth=2, label=f'Median: {median_conf:.3f}')

        ax.set_xlabel('Confidence', fontsize=12, fontweight='bold')
        ax.set_ylabel('Frequency', fontsize=12, fontweight='bold')
        ax.set_title(f'Confidence Distribution (n={len(confidences)})',
                    fontsize=14, fontweight='bold')
        ax.legend(fontsize=10)
        ax.grid(True, alpha=0.3)

        plt.tight_layout()

        img = StatisticsPlotter._fig_to_image(fig)
        plt.close(fig)

        return img

    @staticmethod
    def plot_roi_size_distribution(
        detections: List[Dict],
        figsize: Tuple[int, int] = (10, 6)
    ) -> Image.Image:
        """
        ROI

        Args:
            detections:  (bbox: [x1, y1, x2, y2])
            figsize: 

        Returns:
            PIL.Image
        """
        # 
        widths = []
        heights = []
        for det in detections:
            bbox = det['bbox']
            w = bbox[2] - bbox[0]
            h = bbox[3] - bbox[1]
            widths.append(w)
            heights.append(h)

        if len(widths) == 0:
            fig, ax = plt.subplots(figsize=figsize)
            ax.text(0.5, 0.5, 'No detections',
                   ha='center', va='center', fontsize=16)
            img = StatisticsPlotter._fig_to_image(fig)
            plt.close(fig)
            return img

        # 
        fig, ax = plt.subplots(figsize=figsize)

        # 
        scatter = ax.scatter(widths, heights,
                           c=range(len(widths)),
                           cmap='viridis',
                           alpha=0.6,
                           s=100,
                           edgecolors='black')

        ax.set_xlabel('Width (pixels)', fontsize=12, fontweight='bold')
        ax.set_ylabel('Height (pixels)', fontsize=12, fontweight='bold')
        ax.set_title(f'ROI Size Distribution (n={len(widths)})',
                    fontsize=14, fontweight='bold')
        ax.grid(True, alpha=0.3)

        # 
        cbar = plt.colorbar(scatter, ax=ax)
        cbar.set_label('Detection Index', fontsize=10)

        plt.tight_layout()

        img = StatisticsPlotter._fig_to_image(fig)
        plt.close(fig)

        return img

    @staticmethod
    def plot_iou_distribution(
        ious: List[float],
        bins: int = 20,
        figsize: Tuple[int, int] = (10, 6)
    ) -> Image.Image:
        """
        IoU

        Args:
            ious: IoU
            bins: bin
            figsize: 

        Returns:
            PIL.Image
        """
        if len(ious) == 0:
            fig, ax = plt.subplots(figsize=figsize)
            ax.text(0.5, 0.5, 'No IoU data',
                   ha='center', va='center', fontsize=16)
            img = StatisticsPlotter._fig_to_image(fig)
            plt.close(fig)
            return img

        # 
        fig, ax = plt.subplots(figsize=figsize)

        # 
        n, bins_edges, patches = ax.hist(
            ious,
            bins=bins,
            range=(0, 1),
            color='skyblue',
            alpha=0.7,
            edgecolor='black'
        )

        # 
        mean_iou = np.mean(ious)
        median_iou = np.median(ious)

        ax.axvline(mean_iou, color='red', linestyle='--',
                  linewidth=2, label=f'Mean: {mean_iou:.3f}')
        ax.axvline(median_iou, color='green', linestyle='--',
                  linewidth=2, label=f'Median: {median_iou:.3f}')

        ax.set_xlabel('IoU', fontsize=12, fontweight='bold')
        ax.set_ylabel('Frequency', fontsize=12, fontweight='bold')
        ax.set_title(f'IoU Distribution (n={len(ious)})',
                    fontsize=14, fontweight='bold')
        ax.legend(fontsize=10)
        ax.grid(True, alpha=0.3)

        plt.tight_layout()

        img = StatisticsPlotter._fig_to_image(fig)
        plt.close(fig)

        return img

    @staticmethod
    def plot_comparison(
        before_data: List[float],
        after_data: List[float],
        xlabel: str,
        ylabel: str,
        title: str,
        figsize: Tuple[int, int] = (10, 6)
    ) -> Image.Image:
        """
        Before/After

        Args:
            before_data: 
            after_data: 
            xlabel: x
            ylabel: y
            title: 
            figsize: 

        Returns:
            PIL.Image
        """
        fig, ax = plt.subplots(figsize=figsize)

        x = np.arange(len(before_data))
        width = 0.35

        bars1 = ax.bar(x - width/2, before_data, width,
                      label='Before', color='lightcoral', alpha=0.7, edgecolor='black')
        bars2 = ax.bar(x + width/2, after_data, width,
                      label='After', color='lightgreen', alpha=0.7, edgecolor='black')

        ax.set_xlabel(xlabel, fontsize=12, fontweight='bold')
        ax.set_ylabel(ylabel, fontsize=12, fontweight='bold')
        ax.set_title(title, fontsize=14, fontweight='bold')
        ax.set_xticks(x)
        ax.legend(fontsize=10)
        ax.grid(True, alpha=0.3, axis='y')

        plt.tight_layout()

        img = StatisticsPlotter._fig_to_image(fig)
        plt.close(fig)

        return img

    @staticmethod
    def plot_summary_stats(
        detections: List[Dict],
        class_names: List[str],
        figsize: Tuple[int, int] = (14, 10)
    ) -> Image.Image:
        """
        2x2

        Args:
            detections: 
            class_names: 
            figsize: 

        Returns:
            PIL.Image
        """
        fig = plt.figure(figsize=figsize)
        gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)

        # 1. 
        ax1 = fig.add_subplot(gs[0, 0])
        class_counts = {name: 0 for name in class_names}
        for det in detections:
            cls = det['class']
            if cls in class_counts:
                class_counts[cls] += 1

        classes = list(class_counts.keys())
        counts = list(class_counts.values())
        bars = ax1.bar(classes, counts, color='steelblue', alpha=0.7, edgecolor='black')
        ax1.set_xlabel('Class', fontweight='bold')
        ax1.set_ylabel('Count', fontweight='bold')
        ax1.set_title('Class Distribution', fontweight='bold')
        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45, ha='right')

        # 2. 
        ax2 = fig.add_subplot(gs[0, 1])
        if detections:
            confidences = [det['confidence'] for det in detections]
            ax2.hist(confidences, bins=20, range=(0, 1),
                    color='coral', alpha=0.7, edgecolor='black')
            mean_conf = np.mean(confidences)
            ax2.axvline(mean_conf, color='red', linestyle='--', linewidth=2)
            ax2.set_xlabel('Confidence', fontweight='bold')
            ax2.set_ylabel('Frequency', fontweight='bold')
            ax2.set_title(f'Confidence (μ={mean_conf:.3f})', fontweight='bold')

        # 3. ROI
        ax3 = fig.add_subplot(gs[1, 0])
        if detections:
            widths = [det['bbox'][2] - det['bbox'][0] for det in detections]
            heights = [det['bbox'][3] - det['bbox'][1] for det in detections]
            ax3.scatter(widths, heights, c=range(len(widths)),
                       cmap='viridis', alpha=0.6, s=50, edgecolors='black')
            ax3.set_xlabel('Width (px)', fontweight='bold')
            ax3.set_ylabel('Height (px)', fontweight='bold')
            ax3.set_title('ROI Size Distribution', fontweight='bold')

        # 4. 
        ax4 = fig.add_subplot(gs[1, 1])
        ax4.axis('off')

        stats_text = f"""
         Summary Statistics

        Total Detections: {len(detections)}
        Unique Classes: {sum(1 for c in counts if c > 0)}

        Confidence:
          Mean: {np.mean(confidences):.3f if detections else 'N/A'}
          Std: {np.std(confidences):.3f if detections else 'N/A'}
          Min: {np.min(confidences):.3f if detections else 'N/A'}
          Max: {np.max(confidences):.3f if detections else 'N/A'}

        ROI Size:
          Avg Width: {np.mean(widths):.1f if detections else 'N/A'}
          Avg Height: {np.mean(heights):.1f if detections else 'N/A'}
        """

        ax4.text(0.1, 0.5, stats_text,
                fontsize=10, verticalalignment='center',
                family='monospace',
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

        img = StatisticsPlotter._fig_to_image(fig)
        plt.close(fig)

        return img

    @staticmethod
    def _fig_to_image(fig) -> Image.Image:
        """matplotlib figurePIL Image"""
        fig.canvas.draw()
        img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        return Image.fromarray(img_array)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    print("Statistics Plotter Module")
    print("This module provides various statistical plots")
    print("Import and use in your code:")
    print("  from src.visualization.statistics import StatisticsPlotter")
